﻿using System;
using System.Collections.Generic;
using Allegro.Framework;
using Allegro.MathInterface;

namespace Allegro.Mathlib
{
    public class ConjugateGradientMethod
    {
        // Here ITMAX is the max. allowed number of iterations, while EPS is a small number to
        // rectify the special case of converging to exactly zero function value.
        private const int ITMAX = 500;
        private const double EPS = 1.0e-10;
        private const double TOL = 2.0e-4;

        private BrentFunctionMinimizationMethod _lineMinimizer;
        private double[] _pcom;
        private double[] _xicom;
        private double[] _xt;
        private double[] _g;
        private double[] _h;
        private double[] _xi;
        private int _dimension;  // Vector dimension
        private INumericalVectorDifferentiation _gradient;

        public ConjugateGradientMethod(int dimension, VectorFunction function, VectorMap gradient, double tolerance)
        {
            Tolerance = tolerance;
            _dimension = dimension;
            _lineMinimizer = new BrentFunctionMinimizationMethod(TOL);
            _pcom = new double[_dimension];
            _xicom = new double[_dimension];
            _xt = new double[_dimension];
            _g = new double[_dimension];
            _h = new double[_dimension];
            _xi = new double[_dimension];

            Func = function;
            Df = gradient;
        }
        public ConjugateGradientMethod(int dimension, VectorFunction function, INumericalVectorDifferentiation gradient, double tolerance)
            : this(dimension, function, gradient.Df, tolerance)
        {
            _gradient = gradient;
        }
        public ConjugateGradientMethod(int dimension, VectorFunction function, double tolerance, double[] steps)
            : this(dimension, function, new Gradient(function, steps), tolerance)
        {
        }

        /// <summary>
        /// Vector function to be minimized
        /// </summary>
        public VectorFunction Func { get; private set; }

        /// <summary>
        /// Gradient of the vector function
        /// </summary>
        public VectorMap Df { get; private set; }

        /// <summary>
        /// Number of iterations performed
        /// </summary>
        public double NumIterations;

        /// <summary>
        /// Tolerance for solution
        /// </summary>
        public double Tolerance { get; set; }

        public double FindMinimum(double[] p)
        {
            double fmin = double.NaN;
            int j, its;
            double gg, gam, fp, dgg;
            double g, h, xi;
            double tol = Tolerance;

            if(tol <= 0.0)
            {
                ErrorHandler.Error("Tolerance must be strictly positive");
            }

            NumIterations = 0;

            // Initializations
            fp = Func(p);
            Df(p, _xi);

            for (j = 0; j < _dimension; j++)
            {
                _g[j] = -_xi[j];
                _xi[j] = _h[j] = _g[j];
            }

            // Loop over iterations
            for (its = 0; its < ITMAX; its++)
            {
                NumIterations = its;
                fmin = LineMinimization(p, _xi);
                if (2.0 * Math.Abs(fmin - fp) <= tol * (Math.Abs(fmin) + Math.Abs(fp) + EPS))
                {
                    // Normal return
                    return fmin;
                }
                fp = fmin;
                Df(p, _xi);
                dgg = gg = 0.0;
                for (j = 0; j < _dimension; j++)
                {
                    gg += _g[j] * _g[j];
                    // dgg += _xi[j]*_xi[j] // This statement for Fletcher-Reeves method
                    dgg += (_xi[j] + _g[j]) * _xi[j]; // This statement for Polak-Ribiere
                }
                if (gg == 0.0)
                {
                    // Unlikely. If gradient is exactly zero, then we are already done
                    return fmin;
                }
                gam = dgg / gg;
                for (j = 0; j < _dimension; j++)
                {
                    _g[j] = -_xi[j];
                    _xi[j] = _h[j] = _g[j] + gam * _h[j];
                }
            }
            ErrorHandler.Error("ConjugateGradientMethod.FindMinimum: Too many iterations");
            return fmin;
        }
        /// <summary>
        /// Given an n-dimensional point p[1..n] and an n-dimensional direction xi[1..n], moves and resets p
        /// to where the function Func(p) takes on a minimum along the direction xi from p,
        /// and replaces xi by the actual vector dispacement that p was moved. Also returns the value of func
        /// at the returned location p. This is actually all accomplished by calling the routines FunctionOptimization.BracketMinimum
        /// and BrentFunctionMinimizationMethod.FindMinimum.
        /// </summary>
        /// <param name="p"></param>
        /// <param name="xi"></param>
        /// <returns></returns>
        double LineMinimization(double[] p, double[] xi)
        {
            int j;
            double xx, xmin, fx, fb, fa, bx, ax;
            double fret;
            if (p.Length != _dimension)
            {
                ErrorHandler.Error("ConjugateGradientMethod.LineMinimization: p argument has wrong dimension");
            }
            if (xi.Length != _dimension)
            {
                ErrorHandler.Error("ConjugateGradientMethod.LineMinimization: xi argument has wrong dimension");
            }
            if (Func == null)
            {
                ErrorHandler.Error("ConjugateGradientMethod.LineMinimization: no vector function assigned");
            }

            for (j = 0; j < _dimension; j++)
            {
                _pcom[j] = p[j];
                _xicom[j] = xi[j];
            }
            // Initial guess for brackets
            ax = 0.0;
            xx = 1.0;
            // Bracket minimum
            _lineMinimizer.BracketMinimum(f1dim, ref ax, ref xx, out bx, out fa, out fx, out fb);
            // Minimize
            fret = _lineMinimizer.FindMinimum(f1dim, ax, xx, bx, out xmin);
            // Construct the vector results to return
            for (j = 0; j < _dimension; j++)
            {
                xi[j] *= xmin;
                p[j] += xi[j];
            }

            return fret;
        }

        double f1dim(double x)
        {
            int j;
            double f;

            for (j = 0; j < _dimension; j++)
            {
                _xt[j] = _pcom[j] + x * _xicom[j];
            }
            f = Func(_xt);
            return f;
        }
    }
}
